/*
iovpn
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/songgao/water"
	"github.com/songgao/water/waterutil"
)

type config struct {
	Devtype    string
	Devnetwork string
	Localaddr  string
	Peeraddrs  string
	timestamp  bool
}

type Table struct {
	Ipmap *sync.Map
	Peers []string
	ifio  *water.Interface
	log   func(a ...any)
}

func (t *Table) AddElms(peer ...string) {
	ps := make(map[string]bool, 0)
	tmp := make([]string, 0)
	for _, v := range t.Peers {
		ps[v] = true
	}
	for _, v := range peer {
		ps[v] = true
	}
	for k := range ps {
		tmp = append(tmp, k)
	}
	t.Peers = tmp
}

func (t *Table) DeleteElms(peer ...string) {
	ps := make(map[string]bool, 0)
	tmp := make([]string, 0)
	for _, v := range t.Peers {
		ps[v] = true
	}
	for _, v := range peer {
		delete(ps, v)
	}
	t.Peers = tmp
}

func main() {
	var Conf = &config{}
	flag.BoolVar(&Conf.timestamp, "t", false, "log timestamp")
	flag.StringVar(&Conf.Devtype, "i", "tun", "type of tuntap device[tap|tun] ")
	flag.StringVar(&Conf.Devnetwork, "net", "10.0.0.1/32", "ip address of tun/tap")
	flag.StringVar(&Conf.Peeraddrs, "p", "", "peer addresses used [,] split")
	flag.StringVar(&Conf.Localaddr, "l", "0.0.0.0:8001", "local  address")
	flag.Parse()

	TransUdp(Conf, createtuntap(Conf))
}

func createtuntap(conf *config) *water.Interface {
	fmt.Println(SetColor("1", "32", "["+os.Args[0]+" init]"))
	defer fmt.Println(SetColor("1", "32", "["+os.Args[0]+" start]"))
	// 选择tuntap类型，及设置配置
	config := water.Config{}
	switch conf.Devtype {
	case "tap":
		config.DeviceType = water.TAP
	default:
		config.DeviceType = water.TUN
	}
	// 根据配置生成接口
	ifce, err := water.New(config)
	if err != nil {
		log.Fatal(err)
	}
	if runtime.GOOS == "linux" {
		err = setupNetDevice("ip", "addr", "add", conf.Devnetwork, "dev", ifce.Name())
		if err != nil {
			fmt.Println("[", ifce.Name(), "]:设置网络地址失败：", err)
		} else {
			fmt.Println("[", ifce.Name(), "]:设置地址：", conf.Devnetwork)
		}
		err = setupNetDevice("ip", "link", "set", ifce.Name(), "mtu", "1400", "up")
		if err != nil {
			fmt.Println("[", ifce.Name(), "]:启动失败：", err)
		} else {
			fmt.Println("[", ifce.Name(), "]:启动成功")
		}
	}
	return ifce
}

func setupNetDevice(cmd ...string) error {
	cmdexec := exec.Command(cmd[0], cmd[1:]...)
	cmdexec.Stderr = os.Stderr
	cmdexec.Stdout = os.Stdout
	err := cmdexec.Run()
	if err != nil {
		return err
	}
	return nil
}

/*udp*/
func TransUdp(conf *config, ifio *water.Interface) {
	defer ifio.Close()
	//1.解析公网节点地址
	var tab = &Table{Peers: make([]string, 0), Ipmap: &sync.Map{}, ifio: ifio, log: func() func(a ...any) {
		if conf.timestamp {
			return func(a ...any) { fmt.Println(append([]any{time.Now().Format("2006-01-02 15:04:05.000")}, a...)...) }
		} else {
			return func(a ...any) { fmt.Println(a...) }
		}
	}()}
	if conf.timestamp {
		tab.log = func(a ...any) {
			fmt.Println(append([]any{time.Now().Format("2006-01-02 15:04:05.000")}, a...)...)
		}
	} else {
		tab.log = func(a ...any) {
			fmt.Println(a...)
		}
	}
	tab.Peers = strings.Split(conf.Peeraddrs, ",")
	for _, peer := range tab.Peers {
		paddr, err := net.ResolveUDPAddr("udp", peer)
		if err != nil {
			tab.log(SetColor("3", "31", "parse udp address error:"), err)
			tab.DeleteElms(peer)
			continue
		} else if peer == "" {
			tab.DeleteElms(peer)
			continue
		}
		tab.log(SetColor("2", "32", "[Add peer]: <== "), paddr.String())
	}
	//2.解析本地地址
	localIP, _, err := net.ParseCIDR(conf.Devnetwork)
	if err != nil {
		tab.log(SetColor("3", "31", "parse local address error:"), err)
		return
	}
	//3.监听udp
	locaddr, err := net.ResolveUDPAddr("udp", conf.Localaddr)
	if err != nil {
		tab.log(SetColor("3", "31", "parse udp address error:"), err)
		return
	}
	listener, err := net.ListenUDP("udp", locaddr)
	if err != nil {
		tab.log(SetColor("3", "31", "listen error:"), err)
		return
	}
	defer listener.Close()
	switch {
	case tab.ifio.IsTAP():
		// net -> dev
		go func() {
			buf := make([]byte, 1500)
			for {
				n, clientaddr, err := listener.ReadFromUDP(buf)
				if err != nil {
					tab.log(SetColor("3", "31", "resv udp error:"), err)
					continue
				}
				// keepalive
				if ip, ok := KeepAlive(buf[:n]); ok {
					_, ok := tab.Ipmap.Load(ip.String())
					if !ok {
						tab.log(SetColor("2", "32", "[Ipmap]add:"), ip.String(), "==>", clientaddr.String())
					}
					tab.Ipmap.Store(ip.String(), clientaddr.String())
					continue
				}
				dstMAC := waterutil.MACDestination(buf[:n])
				switch waterutil.MACEthertype(buf[:n]) {
				case waterutil.ARP:
					ARPsIP := ARPsrcIP(buf[:n])
					ARPdIP := ARPdstIP(buf[:n])
					if waterutil.IsBroadcast(dstMAC) {
						v, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), clientaddr.String())
						if !ok {
							log.Println(SetColor("2", "32", "[Ipmap]:add"), ARPsIP.String(), "==>", clientaddr.String())
						} else {
							if v.(string) != clientaddr.String() {
								tab.Ipmap.Store(ARPsIP.String(), clientaddr.String())
							}
						}
						if ARPdIP.String() != localIP.String() {
							v, ok := tab.Ipmap.Load(ARPdIP.String())
							if !ok {
								tab.log(SetColor("3", "31", "not found ip in map:"), ARPdIP.String())
								continue
							} else {
								remoteaddr, err := net.ResolveUDPAddr("udp", v.(string))
								if err != nil {
									tab.log(SetColor("3", "31", "transfor parse udp address error:"), err)
									continue
								}
								_, err = listener.WriteToUDP(buf[:n], remoteaddr)
								if err != nil {
									tab.log(SetColor("3", "31", "transfor upd error:"), err)
									continue
								}
							}
						} else {
							tab.ifio.Write(buf[:n])
						}
					} else {
						v, ok := tab.Ipmap.LoadOrStore(ARPsIP.String(), clientaddr.String())
						if !ok {
							tab.log(SetColor("2", "32", "[Ipmap]:add"), ARPsIP.String(), "==>", clientaddr.String())
						} else {
							if v.(string) != clientaddr.String() {
								tab.log(ARPsIP.String(), clientaddr.String())
							}
						}
						if ARPdIP.String() != localIP.String() {
							v, ok := tab.Ipmap.Load(ARPdIP.String())
							if !ok {
								tab.log(SetColor("3", "31", "not found ip in map:"), ARPdIP.String())
								continue
							} else {
								remoteaddr, err := net.ResolveUDPAddr("udp", v.(string))
								if err != nil {
									tab.log(SetColor("3", "31", "transfor parse udp address error:"), err)
									continue
								}
								_, err = listener.WriteToUDP(buf[:n], remoteaddr)
								if err != nil {
									tab.log(SetColor("3", "31", "transfor upd error:"), err)
									continue
								}
							}
						} else {
							tab.ifio.Write(buf[:n])
						}
					}
				case waterutil.IPv4:
					dstIP := waterutil.IPv4Destination(waterutil.MACPayload(buf[:n]))
					if dstIP.String() != localIP.String() {
						v, ok := tab.Ipmap.Load(dstIP.String())
						if !ok {
							tab.log(SetColor("3", "31", "not found ip in map:"), dstIP.String())
							continue
						} else {
							remoteaddr, err := net.ResolveUDPAddr("udp", v.(string))
							if err != nil {
								tab.log(SetColor("3", "31", "transfor parse udp address error:"), err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], remoteaddr)
							if err != nil {
								tab.log(SetColor("3", "31", "transfor upd error:"), err)
								continue
							}
						}
					} else {
						tab.ifio.Write(buf[:n])
					}
				}
			}
		}()
		// dev -> net
		func() {
			// udp心跳
			go func() {
				for {
					if len(tab.Peers) > 0 {
						for _, peer := range tab.Peers {
							paddr, err := net.ResolveUDPAddr("udp", peer)
							if err != nil {
								tab.log(SetColor("3", "31", "parse udp address error:"), err)
								continue
							}
							listener.WriteToUDP(localIP, paddr)
						}
					}
					time.Sleep(5 * time.Second)
				}
			}()
			buf := make([]byte, 1500)
			for {
				n, err := ifio.Read(buf)
				if err != nil {
					tab.log(err)
					continue
				}
				//	srcMAC := waterutil.MACSource(buf[:n])
				dstMAC := waterutil.MACDestination(buf[:n])
				switch waterutil.MACEthertype(buf[:n]) {
				case waterutil.ARP:
					//ARPsIP := ARPsrcIP(buf[:n])
					ARPdIP := ARPdstIP(buf[:n])
					if waterutil.IsBroadcast(dstMAC) {
						v, ok := tab.Ipmap.Load(ARPdIP.String())
						if !ok {
							if len(tab.Peers) > 0 {
								for _, peer := range tab.Peers {
									paddr, err := net.ResolveUDPAddr("udp", peer)
									if err != nil {
										tab.log(err)
										continue
									}
									_, err = listener.WriteToUDP(buf[:n], paddr)
									if err != nil {
										tab.log(SetColor("3", "31", "send upd error:"), err)
										continue
									}
								}
							}
						} else {
							paddr, err := net.ResolveUDPAddr("udp", v.(string))
							if err != nil {
								tab.log(err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], paddr)
							if err != nil {
								tab.log(SetColor("3", "31", "send upd error:"), err)
								continue
							}
						}
					} else {
						//fmt.Println("[ARP]resp src:", srcMAC.String(), "dst:", dstMAC.String(), "resp ip:", ARPsIP.String())
						raddr, ok := tab.Ipmap.Load(ARPdIP.String())
						if ok {
							paddr, err := net.ResolveUDPAddr("udp", raddr.(string))
							if err != nil {
								tab.log(SetColor("3", "31", "parse udp address error:"), err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], paddr)
							if err != nil {
								tab.log(SetColor("3", "31", "send udp error:"), err)
								continue
							}
						} else {
							tab.log(SetColor("3", "31", "found Ipmap error"))
							continue
						}
					}
				case waterutil.IPv4:
					//fmt.Println("[IPv4]src:", srcMAC.String(), "-->", "dst:", dstMAC.String())
					//fmt.Println("[IPv4]dest ip:", waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
					v, ok := tab.Ipmap.Load(waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
					if !ok {
						tab.log(SetColor("3", "31", "not found ip in map:"), waterutil.IPv4Destination(waterutil.MACPayload(buf[:n])).String())
						continue
					} else {
						remoteaddr, err := net.ResolveUDPAddr("udp", v.(string))
						if err != nil {
							tab.log(SetColor("3", "31", "parse udp address error:"), err)
							continue
						}
						_, err = listener.WriteToUDP(buf[:n], remoteaddr)
						if err != nil {
							tab.log(SetColor("3", "31", "send upd error:"), err)
							continue
						}
					}
				}
			}
		}()
	case ifio.IsTUN():
		// net -> dev
		go func() {
			buf := make([]byte, 1500)
			for {
				n, clientaddr, err := listener.ReadFromUDP(buf)
				if err != nil {
					tab.log(SetColor("3", "31", "resv udp error:"), err)
					continue
				}
				// keepalive
				if ip, ok := KeepAlive(buf[:n]); ok {
					_, ok := tab.Ipmap.Load(ip.String())
					if !ok {
						tab.log(SetColor("2", "32", "[Ipmap]add:"), ip.String(), "==>", clientaddr.String())
					}
					tab.Ipmap.Store(ip.String(), clientaddr.String())
					continue
				}

				if waterutil.IsIPv4(buf[:n]) {
					srcIP := waterutil.IPv4Source(buf[:n])
					dstIP := waterutil.IPv4Destination(buf[:n])
					// 查ipmap是否有src ipmap
					v, ok := tab.Ipmap.LoadOrStore(srcIP.String(), clientaddr.String())
					if !ok {
						tab.log(SetColor("2", "32", "[Ipmap]:add"), srcIP.String(), "==>", clientaddr.String())
					} else {
						if v.(string) != clientaddr.String() {
							tab.Ipmap.Store(srcIP.String(), clientaddr.String())
						}
					}
					// 查ipmap是否到本地,否则转发
					if dstIP.String() != localIP.String() {
						v, ok := tab.Ipmap.Load(dstIP.String())
						if !ok {
							tab.log(SetColor("3", "31", "not found ip in map:"), dstIP.String())
							continue
						} else {
							remoteaddr, err := net.ResolveUDPAddr("udp", v.(string))
							if err != nil {
								tab.log(SetColor("3", "31", "transfor parse udp address error:"), err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], remoteaddr)
							if err != nil {
								tab.log(SetColor("3", "31", "transfor upd error:"), err)
								continue
							}
						}
					} else {
						tab.ifio.Write(buf[:n])
					}
				}
			}
		}()
		// device --> udp
		func() {
			// udp心跳
			go func() {
				for {
					if len(tab.Peers) > 0 {
						for _, peer := range tab.Peers {
							paddr, err := net.ResolveUDPAddr("udp", peer)
							if err != nil {
								tab.log(SetColor("3", "31", "parse udp address error:"), err)
								continue
							}
							listener.WriteToUDP(localIP, paddr)
						}
					}
					time.Sleep(5 * time.Second)
				}
			}()
			buf := make([]byte, 1500)
			for {
				n, err := tab.ifio.Read(buf)
				if err != nil {
					tab.log(err)
					continue
				}
				dstIP := waterutil.IPv4Destination(buf)
				v, ok := tab.Ipmap.Load(dstIP.String())
				if ok {
					remote, err := net.ResolveUDPAddr("udp", v.(string))
					if err != nil {
						tab.Ipmap.Delete(dstIP.String())
						tab.log(SetColor("3", "31", "[IPmap]:Delete"), "==> ", v.(string), err)
						continue
					}
					_, err = listener.WriteToUDP(buf[:n], remote)
					if err != nil {
						tab.log(err)
						continue
					}
				} else {
					if len(tab.Peers) < 1 {
						tab.Ipmap.Range(func(key, value any) bool {
							remote, err := net.ResolveUDPAddr("udp", value.(string))
							if err != nil {
								tab.Ipmap.Delete(key)
								tab.log(SetColor("3", "31", "[IPmap]:Delete"), "==> ", value.(string), err)
							}
							_, err = listener.WriteToUDP(buf[:n], remote)
							if err != nil {
								tab.Ipmap.Delete(key)
								tab.log(SetColor("2", "31", "[IPmap]Delete and send udp error: "), "==> ", value.(string), err)
								return true
							}
							return true
						})
					} else {
						for _, peer := range tab.Peers {
							remote, err := net.ResolveUDPAddr("udp", peer)
							if err != nil {
								tab.log(SetColor("3", "31", "parse upd address error:"), err)
								continue
							}
							_, err = listener.WriteToUDP(buf[:n], remote)
							if err != nil {
								tab.log(SetColor("3", "31", "send upd error:"), err)
								continue
							}
						}
					}
				}
			}
		}()
	}
}

// color tools
func SetColor(i, c, in string) string {
	return fmt.Sprint("\033[", i, ";", c, "m", in, "\033[0m")
}

// ARP工具
func ARPsrcIP(frame []byte) net.IP {
	return net.IPv4(frame[28], frame[29], frame[30], frame[31])
}

func ARPdstIP(frame []byte) net.IP {
	return net.IPv4(frame[38], frame[39], frame[40], frame[41])
}

func KeepAlive(buf []byte) (net.IP, bool) {
	if len(buf) != 16 {
		return nil, false
	} else {
		return net.IPv4(buf[12], buf[13], buf[14], buf[15]), true
	}
}

// 日志

//[255 255 255 255 255 255 34 210 66 114 118 138 8 6 0 1 8 0 6 4 0 1 34 210 66 114 118 138 10 0 0 4 0 0 0 0 0 0 10 0 0 9]
//                  dstmac                srcmac12131415161718192021                srcmac 28                   38
